# _*_ coding : utf-8 _*_
# @Time : 2024-01-30 0:02
# @Author : haowen
# @File : utils
# @Project : pneumonia_tools

import cv2
import math
import cc3d
import numpy as np


# 灰度线性拉伸
def stretch_linear(img_arr):
    return 255.0 * (img_arr - np.min(img_arr)) / (np.max(img_arr) - np.min(img_arr))


# 获取2D凸包络区域
def calc_convex_hull_2d(img_arr):
    # 1.获取输入区域所有像素坐标
    y_arr, x_arr = np.where(img_arr > 127)
    regions = np.concatenate((np.expand_dims(np.expand_dims(x_arr, axis=-1), axis=-1),
                              np.expand_dims(np.expand_dims(y_arr, axis=-1), axis=-1)), axis=-1).astype(np.int32)

    # 2.计算凸包络
    hull_points = cv2.convexHull(regions)

    # 3.填充凸包络
    hull_arr = np.zeros_like(img_arr, dtype=np.uint8)
    cv2.fillPoly(hull_arr, [hull_points], 255)

    return hull_arr


# 根据均值mu_arr向量和协方差矩阵cov_arr构建2D高斯热力图
def put_heatmap_2d(heatmap, mu_arr, cov_arr):
    assert heatmap.ndim == 2
    assert mu_arr.shape == (2, 1) and cov_arr.shape == (2, 2)
    assert (cov_arr == cov_arr.T).all()

    th = 4.6052
    delta = math.sqrt(th * 2)  # 3.0348640826238
    if np.linalg.det(cov_arr) == 0:  # 确保协方差矩阵cov_arr正定
        cov_arr += 1e-12 * np.eye(cov_arr.shape[0], cov_arr.shape[1])
    coinv_arr = np.linalg.inv(cov_arr)  # 求协方差矩阵cov_arr的逆矩阵

    # 限定计算热力图分布坐标范围
    start_coord = mu_arr.squeeze(-1) - delta * np.sqrt(np.max(cov_arr, axis=-1))
    start_coord[start_coord < 0] = 0
    start_coord = start_coord.astype(np.int32)

    end_coord = mu_arr.squeeze(-1) + delta * np.sqrt(np.max(cov_arr, axis=-1))
    end_coord = np.minimum(end_coord, np.array(heatmap.shape[::-1]))
    end_coord = end_coord.astype(np.int32)

    pp = np.mgrid[start_coord[0]:end_coord[0]:1, start_coord[1]:end_coord[1]:1]
    pp = np.reshape(pp, (2, -1))
    for idx in range(pp.shape[-1]):  # 遍历限定计算热力图分布范围的点计算热力值
        exp = np.dot(np.dot((pp[:, idx:idx + 1] - mu_arr).T, coinv_arr), pp[:, idx:idx + 1] - mu_arr).squeeze() / 2.0
        if exp > th:
            continue
        heatmap[pp[1, idx], pp[0, idx]] = max(heatmap[pp[1, idx], pp[0, idx]], math.exp(-exp))
        heatmap[pp[1, idx], pp[0, idx]] = min(heatmap[pp[1, idx], pp[0, idx]], 1.0)

    return heatmap

# 获取3D连通域信息
def calc_cc3d(img_arr):
    labels_out, N = cc3d.connected_components(img_arr, connectivity=26, return_N=True)
    ccinfo_dct = {"S": [], "mean": [], "cov": []}
    for idx in range(1, N + 1):
        coord_arr = np.array(np.where(labels_out == idx), dtype=np.int32)[::-1, :].T
        ccinfo_dct["S"].append((labels_out == idx).sum())
        ccinfo_dct["mean"].append(np.expand_dims(np.mean(coord_arr, axis=0), axis=-1))
        ccinfo_dct["cov"].append(np.cov(coord_arr, rowvar=False))
    ccinfo_dct["S"].append(np.mean(np.array(ccinfo_dct["S"]), axis=0))
    ccinfo_dct["cov"].append(np.cov(np.array(ccinfo_dct["mean"]).squeeze(-1), rowvar=False))
    ccinfo_dct["mean"].append(np.mean(np.array(ccinfo_dct["mean"]), axis=0))

    return ccinfo_dct
